from langchain.schema import Document
from sklearn.preprocessing import MinMaxScaler
from typing import List, Optional
import numpy as np

class HybridRetriever:
    """
    混合检索器，统一调度向量库与BM25检索，并进行重排序
    """
    def __init__(self,
                 vector_retriever,
                 bm25_retriever=None,
                 reranker=None,
                 alpha: float = 0.5,
                 top_k: int = 5):
        """
        初始化混合检索器
        
        参数:
            vector_retriever: 向量检索器，通常是vectorstore.as_retriever()
            bm25_retriever: BM25检索器，可选
            reranker: 重排序模型，如bge-reranker-v2-m3
            alpha: BM25得分与向量得分混合权重，默认0.5
            top_k: 最终返回的Top K块
        """
        self.vector_retriever = vector_retriever
        self.bm25_retriever = bm25_retriever
        self.reranker = reranker
        self.alpha = alpha
        self.top_k = top_k
        # 添加属性以存储各个检索器的结果，便于记录日志
        self.vector_docs = []
        self.bm25_docs = []
        self.hybrid_scores = []

    def normalize(self, scores: List[float]) -> List[float]:
        """
        归一化得分
        """
        if not scores:
            return []
            
        scores_reshaped = np.array(scores).reshape(-1, 1)
        return MinMaxScaler().fit_transform(scores_reshaped).flatten().tolist()

    def invoke(self, query: str) -> List[Document]:
        """
        执行混合检索并返回最终 rerank 后的文档
        """
        # 从向量检索器获取候选
        self.vector_docs = self.vector_retriever.invoke(query)
        
        # 如果有BM25检索器，获取BM25结果
        self.bm25_docs = []
        if self.bm25_retriever:
            self.bm25_docs = self.bm25_retriever.invoke(query)

        # 合并去重
        candidates = {doc.page_content: doc for doc in (self.vector_docs + self.bm25_docs)}
        texts = list(candidates.keys())

        if not texts:
            return []

        # 分别计算各自的得分并归一化
        vector_scores = [1.0 if doc in self.vector_docs else 0.0 for doc in candidates.values()]
        bm25_scores = [1.0 if doc in self.bm25_docs else 0.0 for doc in candidates.values()]

        # 先各自归一化
        if self.bm25_retriever:
            vector_norm = self.normalize(vector_scores)
            bm25_norm = self.normalize(bm25_scores)
            
            # 混合得分
            self.hybrid_scores = [self.alpha * b + (1 - self.alpha) * v for b, v in zip(bm25_norm, vector_norm)]
        else:
            # 如果没有BM25检索器，直接使用向量得分
            self.hybrid_scores = vector_scores

        # 重新组织成text-score对
        text_score_pairs = list(zip(texts, self.hybrid_scores))

        # 按混合得分排序
        text_score_pairs.sort(key=lambda x: x[1], reverse=True)

        # 取前top_k*2个文档进行重排序
        sorted_texts = [text for text, _ in text_score_pairs[:self.top_k * 2]]

        # 如果有reranker，进行重排序
        if self.reranker and len(sorted_texts) > 1:
            rerank_scores = self.reranker.compute_score([(query, text) for text in sorted_texts])
            rerank_indices = sorted(range(len(rerank_scores)), key=lambda i: rerank_scores[i], reverse=True)
            final_texts = [sorted_texts[i] for i in rerank_indices[:self.top_k]]
        else:
            # 无需重排序，直接取前top_k个
            final_texts = sorted_texts[:self.top_k]

        # 返回最终文档
        final_documents = [candidates[text] for text in final_texts]
        return final_documents
        
    def retrieve(self, query: str) -> List[Document]:
        """
        为保持向后兼容性的旧接口，现在调用invoke方法
        """
        import warnings
        warnings.warn(
            "The method `retrieve` was deprecated and will be removed in future. Use `invoke` instead.", 
            DeprecationWarning,
            stacklevel=2
        )
        return self.invoke(query) 